import os
from qagen.utils import *
from random import *
import bpy
import mathutils
from mathutils import Euler, Matrix, Vector
import math
import shutil
from PIL import Image
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
import multiprocessing

cate_list = read_category_list()
obj_list = read_obj_info()

def get_table():
    table_obj = None
    for obj in bpy.data.objects:
        if "TableDining" in obj.name:
            table_obj = obj
            break
    return table_obj
def get_table_():
    table_obj = None
    for obj in bpy.data.objects:
        if "TableDining" in obj.name:
            table_obj = obj
    return table_obj

def get_lamp():
    table_obj = None
    for obj in bpy.data.objects:
        if "PointLamp" in obj.name:
            table_obj = obj
            break
    return table_obj

def get_object_world_bbox(obj):
    mesh = obj.data
    world_verts = [obj.matrix_world @ v.co for v in mesh.vertices]
    # print(len(mesh.vertices))
    return world_verts

def getBbox(parent_obj):
    all_points = []
    bpy.context.view_layer.update()
    if 'TableDining' in parent_obj.name:
        mesh = parent_obj.data
        worlds_verts = [Euler(parent_obj.rotation_euler, parent_obj.rotation_mode).to_matrix().to_4x4() @ v.co for v in mesh.vertices]
        all_points.extend(worlds_verts)
    for obj in parent_obj.children_recursive:
        if obj.type == 'MESH' and not obj.hide_viewport and not obj.hide_render:
            part_bbox = get_object_world_bbox(obj)
            all_points.extend(part_bbox)
    min_corner = Vector((
        min(p.x for p in all_points),
        min(p.y for p in all_points),
        min(p.z for p in all_points)
    ))
    max_corner = Vector((
        max(p.x for p in all_points),
        max(p.y for p in all_points),
        max(p.z for p in all_points)
    ))

    return min_corner, max_corner

def fixLocation(obj, current_center, target_center):
    # global depsgraph
    for obj in obj.children_recursive:
        if obj.children == ():
            obj.location = obj.location - current_center + target_center
            # depsgraph = bpy.context.evaluated_depsgraph_get()

def random_position(sz, tablemin, tablemax):
    tx_min = tablemin.x + sz.x / 2
    tx_max = tablemax.x - sz.x / 2
    ty_min = tablemin.y + sz.y / 2
    ty_max = tablemax.y - sz.y / 2
    tx = uniform(tx_min, tx_max)
    ty = uniform(ty_min, ty_max)
    tz = tablemax.z
    return (tx, ty, tz)

ITV = 0.15
def checkOverlap(obj_nbbox, x, y, sz):
    bxmin = x - sz.x / 2
    bxmax = x + sz.x / 2
    bymin = y - sz.y / 2
    bymax = y + sz.y / 2
    for a_min, a_max in obj_nbbox:
        axmin, aymin, _ = a_min
        axmax, aymax, _ = a_max
        if axmin >= bxmax + ITV or axmax + ITV <= bxmin or aymin >= bymax + ITV or aymax + ITV <= bymin:
            continue
        return True
    return False

def recursive_select(obj):
    obj.select_set(True)
    for child in obj.children_recursive:
        child.select_set(True)

def delete_sketchfab_models(blend_file_path):
    bpy.ops.object.select_all(action='DESELECT')
    for obj in bpy.data.objects:
        if "Sketchfab_model" in obj.name:
            recursive_select(obj)
    bpy.ops.object.delete()
    bpy.ops.wm.save_mainfile(filepath=blend_file_path)

def delete_pillars(blend_file_path):
    bpy.ops.object.select_all(action='DESELECT')
    for obj in bpy.data.objects:
        if "PillarFactory" in obj.name:
            recursive_select(obj)
    bpy.ops.object.delete()
    bpy.ops.wm.save_mainfile(filepath=blend_file_path)

def fix_table(blend_file_path):
    bpy.ops.object.select_all(action='DESELECT')
    table_obj = get_table_()
    table_obj.select_set(True)
    table_obj.rotation_euler.x = 0
    table_obj.rotation_euler.y = 0
    table_obj = get_table()
    table_obj.select_set(True)
    table_obj.rotation_euler.x = 0
    table_obj.rotation_euler.y = 0
    bpy.ops.wm.save_mainfile(filepath=blend_file_path)

def random_camera(scene_path):
    table_obj = get_table()
    if table_obj is None:
        print("未找到名称包含 'TableDining' 的物体！")
    else:
        print(table_obj.location)
        min_corner, max_corner = getBbox(table_obj)
        print(min_corner, max_corner)
        center = (min_corner + max_corner) / 2.0 + table_obj.location # HERE
        width = max_corner.x - min_corner.x
        depth = max_corner.y - min_corner.y
        top_z = max_corner.z

        camera_obj = None
        camrig = None
        for obj in bpy.data.objects:
            if obj.type == 'CAMERA':
                camera_obj = obj
                print(obj.name, obj.type, obj.data.angle_x, obj.data.angle_y)
            if obj.name == 'camrig.0':
                camrig = obj
                print(obj.location, obj.rotation_euler)

        cam_data = camera_obj.data

        fov_x = cam_data.angle_x
        fov_y = cam_data.angle_y

        req_distance_x = (width / 2.0) / math.tan(fov_x / 2.0)
        req_distance_y = (depth / 2.0) / math.tan(fov_y / 2.0)
        req_distance = max(req_distance_x, req_distance_y)
        print("计算所需的距离: %.3f" % req_distance)

        rand_ang = randint(45, 50)
        # req_distance += math.log2(width / depth) * 0.3
        req_distance += (width - depth) * depth * 0.25
        camrig.location = (center.x + req_distance * math.sin(math.radians(rand_ang)), center.y, top_z + req_distance - 0.1)
        camrig.rotation_euler = (math.radians(rand_ang), 0, math.radians(90))
        print("摄像机新位置:", camrig.location)
        print("摄像机新旋转 (度):",
              (math.degrees(camrig.rotation_euler.x),
               math.degrees(camrig.rotation_euler.y),
               math.degrees(camrig.rotation_euler.z)))

        bpy.ops.wm.save_as_mainfile(filepath=scene_path)
        return camrig.location, camrig.rotation_euler, fov_x, fov_y

def random_len(scl):
    scl = max(scl, 10)
    scl = min(scl, 50)
    scl = scl * 0.01 * uniform(0.8, 1.2)
    return scl

def sample_x_coordinate(index, total, table_min, table_max, sz):
    if total <= 1:
        weight = random()
    else:
        # beta_param = 4 - 3 * (index / (total - 1))
        weight = betavariate(index + 1, total - index)
    x = table_min.x + sz.x / 2 + weight * (table_max.x - table_min.x - sz.x)
    return x

def sample_y_coordinate(table_min, table_max, sz, dir):
    if dir.y == 1:
        return table_min.y + sz.y / 2 + random() * (table_max.y - table_min.y - sz.y) / 2
    elif dir.y == -1:
        return table_max.y - sz.y / 2 - random() * (table_max.y - table_min.y - sz.y) / 2
    return table_min.y + sz.y / 2 + random() * (table_max.y - table_min.y - sz.y)

def genqa(id, scene_id):
    scene_path = get_scene_path(scene_id)

    bpy.ops.wm.open_mainfile(filepath=scene_path)

    delete_sketchfab_models(scene_path)
    delete_pillars(scene_path)
    fix_table(scene_path)

    cam_loc, cam_rot, cam_x, cam_y = random_camera(scene_path)

    table = get_table()
    table_min, table_max = getBbox(table)
    table_sz = table_max - table_min
    table_min += table.location
    table_max += table.location
    print(f'{table_min=}, {table_max=}')
    table_max.x -= table_sz.x * 0.10
    table_min.y += table_sz.y * 0.10
    table_max.y -= table_sz.y * 0.10

    while True:
        bpy.ops.wm.open_mainfile(filepath=scene_path)

        while True:
            main_cate = choice(list(cate_list.keys()))
            if len(cate_list[main_cate]) >= 3 and obj_list[cate_list[main_cate][0]]['scale'] <= 35.0:
                break
        
        if randint(1, 2) > 1:
            objs_id = choices(cate_list[main_cate], k = randint(2, 3))
        else:
            objs_id = []

        objs_id.extend(choices(list(obj_list.keys()), k = randint(1, 3) + (3 - len(objs_id))))

        # print(objs_id)
        # print(len(objs_id))

        obj_info = {}
        for obj_id in objs_id:
            info = obj_list[obj_id].copy()
            info['scale'] = random_len(info['scale'])
            bpy.ops.import_scene.gltf(filepath=f'myobj/{info["path"]}')
            tmp_list = bpy.context.selected_objects

            armature = None
            for obj in tmp_list:
                if obj.type == 'ARMATURE':
                    armature = obj
                    break
            if armature != None:
                obj_info = None
                break

            obj = tmp_list[0]
            info['object'] = obj
            old_dir = info['direction']
            info['direction'] = mathutils.Vector((old_dir[1], -old_dir[0], old_dir[2]))

            # set scale
            omin, omax = getBbox(obj)
            # print(f'original bbox: {omax - omin}')
            obj.scale *= info['scale'] / max(omax - omin)
            bpy.context.view_layer.update()

            # set rotation
            if abs(old_dir[2]) == 1:
                ang = uniform(0, 2 * math.pi)
                # obj.rotation_euler[2] = ang
                q = mathutils.Quaternion((0, 0, 1), ang)
                obj.rotation_quaternion = q @ obj.rotation_quaternion
            else:
                omin, omax = getBbox(obj)
                from_dir = info['direction']
                sz = omax - omin
                if (omax.y - omin.y < 0.05 and from_dir.x != 0) or (omax.x - omin.x < 0.05 and from_dir.y != 0):
                    target = choice([mathutils.Vector((0, 1, 0)), mathutils.Vector((0, -1, 0))])
                if (omax.x - omin.x < 0.05 and from_dir.x != 0) or (omax.y - omin.y < 0.05 and from_dir.y != 0):
                    target = mathutils.Vector((1, 0, 0))
                else:
                    target = choice([mathutils.Vector((1, 0, 0)), mathutils.Vector((0, 1, 0)), mathutils.Vector((0, -1, 0))])
                # print(target)
                angle_from = math.atan2(from_dir.y, from_dir.x)
                angle_to = math.atan2(target.y, target.x)
                # print(obj.rotation_mode)
                # obj.rotation_euler[2] = angle_to - angle_from
                q = mathutils.Quaternion((0, 0, 1), angle_to - angle_from)
                obj.rotation_quaternion = q @ obj.rotation_quaternion
                info['direction'] = target
            bpy.context.view_layer.update()

            info['bboxmin'], info['bboxmax'] = getBbox(obj)
            # print(f'{obj.location=}')
            # print(info['bboxmax'] - info['bboxmin'])
            # info['bboxmin'] += info['object'].location
            # info['bboxmax'] += info['object'].location
            obj_info[obj_id] = info
        if obj_info == None:
            continue
        obj_info = dict(sorted(obj_info.items(), key=lambda item: item[1]['bboxmax'].z - item[1]['bboxmin'].z, reverse=True))

        try_times = 0
        while try_times < 20:
            try_times += 1
            obj_nbbox = []
            for idx, (obj_id, info) in enumerate(obj_info.items()):
                obj = info['object']
                obj_sz = info['bboxmax'] - info['bboxmin']
                try_times_x = 0
                while try_times_x < 10:
                    try_times_x += 1
                    x = sample_x_coordinate(idx, len(obj_info), table_min, table_max, obj_sz)
                    try_times_y = 0
                    while try_times_y < 10:
                        try_times_y += 1
                        y = sample_y_coordinate(table_min, table_max, obj_sz, info['direction'])
                        if not checkOverlap(obj_nbbox, x, y, obj_sz):
                            break
                        y = None
                    if y != None:
                        break
                    x = None
                if x == None:
                    obj_nbbox = None
                    break

                # move to (x, y) and append obj_nbbox
                target_co = mathutils.Vector((x, y, table_max.z + obj_sz.z / 2))
                obj_nbbox.append((target_co - obj_sz / 2, target_co + obj_sz / 2))
                source_co = (info['bboxmax'] + info['bboxmin']) / 2
                # print(obj.location, table_max.z, source_co.z, obj_sz.z)
                obj.location = obj.location - source_co + target_co
                bpy.context.view_layer.update()
                info['bboxmin'], info['bboxmax'] = getBbox(obj)
                # info['bboxmin'] += obj.location
                # info['bboxmax'] += obj.location
                # print(obj.location, ((info['bboxmin'] + info['bboxmax']) / 2).z)
                obj_nbbox.append((info['bboxmin'], info['bboxmax']))
            if obj_nbbox != None:
                break

        if obj_nbbox != None:
            break

    bpy.ops.wm.save_as_mainfile()

    render(scene_id)

    # print(f'len={len(obj_info)}')
    # print(obj_info)

    source_image = get_image_path(scene_id)
    source_depth = get_depth_path(scene_id)
    target_image = f'data/{id}.jpg'
    target_depth = f'data/{id}.png'
    with Image.open(source_image) as img:
        if img.mode in ('RGBA', 'LA'):
            img = img.convert('RGB')
        img.save(target_image, 'JPEG')
    data = np.load(source_depth)
    data_min = np.min(data)
    data_max = min(np.max(data), 4)
    data = 255 - 255 * (np.minimum(data, 4) - data_min) / (data_max - data_min)
    data = np.stack((data,) * 3, axis=-1).astype(np.uint8)
    data = Image.fromarray(data)
    data.save(target_depth, 'PNG')

    with open(get_objid_path(scene_id), 'r') as f:
        id_json = json.load(f)
    # read mask.npy
    mask = np.load(get_mask_path(scene_id))
    allinfo = {}
    allinfo['scene_id'] = scene_id
    allinfo['camera'] = {
        'location': [cam_loc.x, cam_loc.y, cam_loc.z],
        'rotation': [cam_rot.x, cam_rot.y, cam_rot.z],
        'fov': [cam_x, cam_y],
    }
    allinfo['table'] = {
        'tablemin': [table_min.x, table_min.y, table_min.z],
        'tablemax': [table_max.x, table_max.y, table_max.z],
    }
    for obj_id, info in obj_info.items():
        info['mask'] = get_object_mask_rle(id_json, mask, info['object'].name)
        info['object'] = None
        info['bboxmin'] = [info['bboxmin'].x, info['bboxmin'].y, info['bboxmin'].z]
        info['bboxmax'] = [info['bboxmax'].x, info['bboxmax'].y, info['bboxmax'].z]
        info['direction'] = [info['direction'].x, info['direction'].y, info['direction'].z]
        obj_info[obj_id] = info
    allinfo['objects'] = obj_info

    with open(f'data/{id}.json', 'w') as f:
        json.dump(allinfo, f, indent=4)

    # delete_sketchfab_models(scene_path)

if __name__ == "__main__":
    with open('Xscene_ids.json', 'r') as f:
        scene_ids = json.load(f)

    # num_processes = 1  # 设置最大并行数
    # num_samples = 1  # 设置要生成的样本总数
    # seed(1)
    # genqa('xxxxxxxxxxxxxxxx', '7')
    # exit(0)

    num_processes = 4
    # num_samples = 8192
    num_samples = 16
    # seed(2333)
    pool = multiprocessing.Pool(processes=num_processes)
    results = []

    for i in range(num_samples):
    # for i in [0]:
        id = random_hex()
        scene_id_to_use = scene_ids[i % len(scene_ids)] # 循环使用预选的场景
        results.append(pool.apply_async(genqa, args=(id, scene_id_to_use)))

    # 等待所有进程完成
    for res in results:
        res.get()

    pool.close()
    pool.join()
